use std::cmp::Ordering;
use std::mem;

/// A binomial heap is a data structure that acts as a priority queue but also supports
/// efficient merging of heaps. This implementation is generic over the type `T`, which
/// must implement the `Ord` trait for comparison.
pub struct BinomialHeap<T> {
    heap_type: HeapType,
    compare_func: fn(&T, &T) -> Ordering,
    num_values: usize,
    roots: Vec<Option<Box<BinomialTree<T>>>>,
}

/// Represents the type of the heap, either a min-heap or a max-heap.
#[derive(Clone, Copy, PartialEq, Debug)]
pub enum HeapType {
    Min,
    Max,
}

/// A binomial tree is a recursive data structure used in the binomial heap.
struct BinomialTree<T> {
    value: T,
    order: u16,
    subtrees: Vec<Box<BinomialTree<T>>>,
}

impl<T> BinomialHeap<T> {
    /// Creates a new binomial heap with the specified heap type and comparison function.
    pub fn new(heap_type: HeapType, compare_func: fn(&T, &T) -> Ordering) -> Self {
        BinomialHeap {
            heap_type,
            compare_func,
            num_values: 0,
            roots: Vec::new(),
        }
    }

    /// Inserts a new value into the heap.
    pub fn insert(&mut self, value: T) {
        let new_tree = Box::new(BinomialTree {
            value,
            order: 0,
            subtrees: Vec::new(),
        });

        let mut fake_heap = BinomialHeap {
            heap_type: self.heap_type,
            compare_func: self.compare_func,
            num_values: 1,
            roots: vec![Some(new_tree)],
        };

        self.merge(&mut fake_heap);
        self.num_values += 1;
    }

    /// Removes and returns the smallest (or largest, depending on heap type) value from the heap.
    pub fn pop(&mut self) -> Option<T> {
        if self.num_values == 0 {
            return None;
        }

        let least_index = self
            .roots
            .iter()
            .enumerate()
            .filter_map(|(i, tree)| tree.as_ref().map(|_| i))
            .min_by(|&a, &b| {
                self.compare(
                    &self.roots[a].as_ref().unwrap().value,
                    &self.roots[b].as_ref().unwrap().value,
                )
            })?;

        let least_tree = self.roots[least_index].take().unwrap();
        let mut fake_heap = BinomialHeap {
            heap_type: self.heap_type,
            compare_func: self.compare_func,
            num_values: least_tree.order as usize,
            roots: least_tree
                .subtrees
                .into_iter()
                .map(|tree| Some(tree))
                .collect(),
        };

        self.merge(&mut fake_heap);
        self.num_values -= 1;

        Some(least_tree.value)
    }

    /// Returns the number of entries in the heap.
    pub fn num_entries(&self) -> usize {
        self.num_values
    }

    /// Merges another heap into this heap.
    fn merge(&mut self, other: &mut BinomialHeap<T>) {
        let max_length = self.roots.len().max(other.roots.len()) + 1;
        let mut new_roots = vec![None; max_length];
        let mut carry = None;

        for i in 0..max_length {
            let mut vals = [
                self.roots.get(i).cloned().unwrap_or(None),
                other.roots.get(i).cloned().unwrap_or(None),
                carry,
            ];

            let num_vals = vals.iter().filter(|&&x| x.is_some()).count();

            if num_vals & 1 != 0 {
                new_roots[i] = vals[2].take();
            }

            if num_vals & 2 != 0 {
                carry = Some(self.merge_trees(vals[0].take().unwrap(), vals[1].take().unwrap()));
            } else {
                carry = None;
            }
        }

        self.roots = new_roots;
        self.num_values += other.num_values;
        other.roots.clear();
    }

    /// Merges two binomial trees of the same order.
    fn merge_trees(&self, mut tree1: Box<BinomialTree<T>>, mut tree2: Box<BinomialTree<T>>) -> Box<BinomialTree<T>> {
        if self.compare(&tree1.value, &tree2.value) == Ordering::Greater {
            mem::swap(&mut tree1, &mut tree2);
        }

        tree1.subtrees.push(tree2);
        tree1.order += 1;
        tree1
    }

    /// Compares two values based on the heap type.
    fn compare(&self, data1: &T, data2: &T) -> Ordering {
        let result = (self.compare_func)(data1, data2);
        if self.heap_type == HeapType::Min {
            result
        } else {
            result.reverse()
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_insert_and_pop() {
        let mut heap = BinomialHeap::new(HeapType::Min, |a, b| a.cmp(b));
        heap.insert(5);
        heap.insert(3);
        heap.insert(8);

        assert_eq!(heap.pop(), Some(3));
        assert_eq!(heap.pop(), Some(5));
        assert_eq!(heap.pop(), Some(8));
        assert_eq!(heap.pop(), None);
    }

    #[test]
    fn test_num_entries() {
        let mut heap = BinomialHeap::new(HeapType::Min, |a, b| a.cmp(b));
        assert_eq!(heap.num_entries(), 0);

        heap.insert(10);
        assert_eq!(heap.num_entries(), 1);

        heap.insert(20);
        assert_eq!(heap.num_entries(), 2);

        heap.pop();
        assert_eq!(heap.num_entries(), 1);
    }
}
